#!/bin/sh -
"exec" "python" "-O" "$0" "$@"

__doc__ = """Bababa lalala bababa lalala.
"""

__version__ = "0.6.4"

_MAX_REQUEST_CACHE_SIZE_ = 1024*1024
_MAX_RESPONSE_CACHE_SIZE_ = 1024*1024

import StringIO, httplib, BaseHTTPServer, select, socket, SocketServer, urlparse

  
class  HttpConnectError(Exception):
    def __init__(self):
        pass
class  HttpClientError(Exception):
    def __init__(self):
        pass
class  HttpServerError(Exception):
    def __init__(self):
        pass
class  HttpProxyError(Exception):
    def __init__(self):
        pass
class  HttpUnknownError(Exception):pass

class ISocket:
    def set_fd(self,fd):pass
    def get_fd(self):pass

class IReadable:
    def read(self,size=4096):pass        

class IWriteable:
    def flush(self):pass
    def write(self,buf):pass

class SocketJoiner:
    defaultMaxIdleCount = 20
    defaultInteval = 3
    def __init__(self,sockA=None,sockB=None,maxIdle=defaultMaxIdleCount):
        self.socks = [sockA,sockB]
        self.maxIdle = maxIdle
        self.inteval = SocketJoiner.defaultInteval
    
    def set_inteval(self,intv):
        self.inteval = intv
        
    @classmethod
    def join_up(clsDef,sockA,sockB,maxIdleCount = defaultMaxIdleCount, inteval = defaultInteval):
        instant = SocketJoiner(sockA,sockB,maxIdleCount)
        instant.set_inteval(inteval)
        instant.transfer()
        
        
    def join(self,sockA,sockB):
        self.socks = [sockA,sockB]
        
    def transfer(self):
        iw = [self.socks[0].get_fd(), self.socks[1].get_fd()]
        ow = []
        count = 0
        try:
            while 1:
                count += 1
                (ins, _, exs) = select.select(iw, ow, iw, self.inteval)
                if exs: break
                if ins:
                    for i in ins:
                        if i is iw[0]:
                            input = self.socks[0]
                            out = self.socks[1]
                        else:
                            input = self.socks[1]
                            out = self.socks[0]
                        data = input.read(8192)                    
                        if data:
                            out.write(data)
                            count = 0
                if count == self.maxIdle: break
        except:
            print "ERRRRRRRRRRRRRRRR"    
            raise    
        finally:
            return
           

class HttpRequestCache:
    def set_fd(self,fd):pass
    
    def get_fd(self):
        return self.connection
    
    def read(self,size=4096):
        return self.connection.recv(size)
    
    def write_response_header(self,buf):
        self._resp_writed_length = 0
        self.connection.send(buf)
    
    def get_writed_length(self):
        return self._resp_writed_length
    
    def write(self,data):
        print len(data)
        self._resp_writed_length += len(data)
        self.connection.send(data)
    
    def flush(self):pass
        
    def __init__(self,proxyReq):
        #copy header
        self._resp_writed_length = -1
        self.command = proxyReq.command
        self.path = proxyReq.path
        self.request_version = proxyReq.request_version
        self.protocol_version = proxyReq.protocol_version
        self.requestline = proxyReq.requestline        
        self.close_connection = proxyReq.close_connection
        self.headers = proxyReq.headers 
        self.wfile = proxyReq.wfile
        self.proxyReq = proxyReq
        #parse request line
        if self.command == "CONNECT":
            self.netloc = self.path 
            self.scm = "https"
            self.params = ""
            self.file_path = ""
            self.query = ""
            self.fragment = ""
        else:
            (scm, netloc, path, params, query, fragment) = urlparse.urlparse(
                self.path)     
            self.scm = scm 
            self.netloc = netloc
            self.params = params
            self.file_path = path
            self.query = query
            self.fragment = fragment
        
        i = self.netloc.find(':')        
        if i >= 0:
            self.host = self.netloc[:i]
            self.port = int(self.netloc[i+1:])
        else:
            self.host = self.netloc
            if self.command == "CONNECT":
                self.port = int(443)
            else:    
                self.port = int(80)
        
        #read request body       
        self.body_cache = StringIO.StringIO()
        try:
            lenStr = self.headers.get("Content-Length","0")
            self.content_length = int(lenStr)            
        except KeyError:pass
            
        self.connection = proxyReq.connection
        if self.content_length > _MAX_REQUEST_CACHE_SIZE_:
            self._is_complete = 0
            return
        elif self.content_length == 0:
            self._is_complete = 1
        else:
            self._is_complete = 0
            while self.content_length - 8192 > self.body_cache.len:
                self.body_cache.write(self.connection.recv(8192))            
            if self.content_length - self.body_cache.len > 0:    
                self.body_cache.write(self.connection.recv(self.content_length - self.body_cache.len))     
            self._is_complete = 1

class HttpConnectResponse:
    
    def __init__(self,sock):
        self.sock_fd = sock
        
    def get_fd(self):
        return self.sock_fd 
    
    def read(self,count=4096):
        return self.sock_fd.recv(count)
    
    def write(self,buf):
        return self.sock_fd.send(buf)
    
    def flush(self):
        pass

class HTTPResponseCache(httplib.HTTPResponse): 
    def __init__(self,sock):
        self.sock_fd = sock
        httplib.HTTPResponse.__init__(self,sock)
        self.cache = None
         
    def cache_body(self):
        if self.cache == None and self.length < _MAX_RESPONSE_CACHE_SIZE_ and self.length != 0:
            self.cache =self.read()
        return self.cache
    
    def get_fd(self):
        return self.sock_fd 
    
    def read(self,count=4096):
        return self.fp.read(count)
    
    def write(self,buf):
        raise HttpUnknownError()
    
    def flush(self):
        raise HttpUnknownError()
    
class BaseHTTPProxyHandler:
    def log_request(self,resp_code):
        pass
    
    def handle(self, request_cache):
        self.req = request_cache
        try:
            self.do_connect()
            if self.req.command == "CONNECT":
                try:
                    self.log_request(200)
                    self._write_connect_response();
                    self.resp = HttpConnectResponse(self.soc)
                    SocketJoiner.join_up(self.req, self.resp)
                    return
                except:
                    raise HttpUnknownError()
        except:
            raise HttpConnectError()
        
        try:
            self.do_send_request()
            self.do_handle_resp()
        except HttpClientError:
            raise
        except HttpServerError:
            raise
        except:
            raise HttpProxyError()
        #TODO: check which side is closed, and raise a correct exception
        
    def do_connect(self):
        try:            
            self.soc = socket.socket(socket.AF_INET, socket.SOCK_STREAM)   
            host_port = self.req.host, self.req.port
            self.soc.connect(host_port)
        except:
            raise 
        print "\n" "connect to %s:%d" % host_port
    
    def _do_get_request_line(self):
        return "%s %s %s\r\n" % (
            self.req.command,
            urlparse.urlunparse(('', '', self.req.file_path, self.req.params, self.req.query, '')),
            self.req.request_version)
        
    
    def do_send_request(self):
        self.log_request(200)
        self.soc.send(self._do_get_request_line())
        self.req.headers['Connection'] = 'close'
        try:
            del self.req.headers['Proxy-Connection']
        finally:
            pass    
        #TODO: check break-point and continue
        for key_val in self.req.headers.items():
            self.soc.send("%s: %s\r\n" % key_val)                       
        self.soc.send("\r\n")
        self.soc.send(self.req.body_cache.getvalue())
        return True
    
    def do_handle_resp(self):
        try:
            self.resp = HTTPResponseCache(self.soc)
            self.resp.begin()
        except:
            raise HttpServerError()
        
        try:
            if self.req.get_writed_length() == -1:
                #TODO: Send Response
                self.req.write_response_header("HTTP/%d.%d %d %s\r\n" % (self.resp.version/10, self.resp.version%10,self.resp.status,self.resp.reason))
            
                respheaders = self.resp.getheaders()
                #respheaders['Connection'] = 'close'  
                  
                #TODO: check break-point and continue
                for key_val in respheaders:
                    if key_val[0].lower() != "connection":
                        self.req.write_response_header("%s: %s\r\n" % key_val)
                
                self.req.write_response_header("Connection: close\r\n")                      
                self.req.write_response_header("\r\n")
        except:
            raise HttpClientError()
        
        try:
            SocketJoiner.join_up(self.req, self.resp)
        except:
            #TODO: check the value 
            raise 
    
    def _send_error(self,errCode, msg):
        self.req.proxyReq.send_error(errCode,msg)        
    
    def _write_back(self, str):
        self.req.write(str)
    
    def _agent_str(self):
        return "JJProxy %s"% __version__
  
    def _write_connect_response(self):
        self._write_back(self.req.protocol_version + " 200 Connection established\r\n")
        self._write_back("Proxy-agent: %s\r\n" % self._agent_str())
        self._write_back("\r\n")          
     
class CNHandler (BaseHTTPProxyHandler):pass
#    def handle(self, request_cache):
#        return BaseHTTPProxyHandler.handle(self, request_cache)

class ForwardHttpProxyHandler (BaseHTTPProxyHandler):
    def __init__(self,proxyHost="127.0.0.1",proxyPort=8000):
        self.proxyHost = proxyHost
        self.proxyPort = proxyPort
    
    def do_connect(self):
        self.soc = socket.socket(socket.AF_INET, socket.SOCK_STREAM)   
        host_port = self.proxyHost, self.proxyPort
        self.soc.connect(host_port) 
        print "\n" "connect to %s:%d" % host_port
    

    def _do_get_request_line(self):
        print self.req.requestline
        return "%s\r\n"%self.req.requestline

class ForwardSock4ProxyHandler (BaseHTTPProxyHandler):
    def handle(self, request_cache):
        return BaseHTTPProxyHandler.handle(self, request_cache)   

class ForwardSock5ProxyHandler (BaseHTTPProxyHandler):
    def handle(self, request_cache):
        return BaseHTTPProxyHandler.handle(self, request_cache) 
    
class HttpProxyHandlerManager:
    blocked_list = dict()   
    ins = None
    @classmethod
    def get_instance(cls):
        if cls.ins == None:
            cls.ins = HttpProxyHandlerManager();            
        return cls.ins
    
    def register(self,hdlrCls):pass
    '''TODO add lock...'''
        
    def query(self,req,order):
        if HttpProxyHandlerManager.blocked_list.has_key((req.host,req.port)):
            return ForwardHttpProxyHandler()
        else:
            return CNHandler()
        
        if order == 0:
            return CNHandler()
        elif order == 1:
            return ForwardHttpProxyHandler()
        return None    
    
    def confirm(self,req,handler): 
        pass
     
    def mark_blocked(self,req,handler): 
        self.__class__.blocked_list[(req.host,req.port)] = handler.__class__;
        pass
    
        
class ProxyRequestDispatchor (BaseHTTPServer.BaseHTTPRequestHandler):
    __base = BaseHTTPServer.BaseHTTPRequestHandler
    __base_handle = __base.handle

    server_version = "JJProxy.org/" + __version__
    rbufsize = 0                        # self.rfile Be unbuffered

    def handle(self):
        (ip,port) =  self.client_address
        if hasattr(self, 'allowed_clients') and ip not in self.allowed_clients:
            self.raw_requestline = self.rfile.readline()
            if self.parse_request(): self.send_error(403)
        else:
            self.__base_handle()

    def dispatch_request(self):
        
        self.req_cache = HttpRequestCache(self)
        if self.command.lower() != "connect":
            if self.req_cache.scm != 'http' or self.req_cache.fragment or not self.req_cache.netloc:
                self.send_error(400, "bad url %s" % self.path)
                return
        order = 0
        manager = HttpProxyHandlerManager.get_instance()
        while True:
            hdlr = manager.query(self.req_cache, order)
            if hdlr:
                try:
                    hdlr.handle(self.req_cache)
                    manager.confirm(self.req_cache, hdlr) 
                except HttpConnectError:
                    print "HttpConnectError"
                    manager.mark_blocked(self.req_cache, hdlr) 
                    order += 1
                    continue
                except HttpClientError:
                    print "HttpClientError"  
                    return
                except HttpServerError:
                    print "ServerClientError"                    
                    manager.mark_blocked(self.req_cache, hdlr) 
                    order += 1
                    continue
                except HttpProxyError:
                    print "HttpProxyError"  
                    pass
                except :
                    print "unknown exception"
                    raise
                finally:
                    print "OK? Right?"   
                
                return                 
            else:
                self.send_error(581, "Fuck GFW!!! Sorry we can't find out a way to %s ..." % self.path)
                return 
            order += 1

    do_CONNECT  = dispatch_request
    do_HEAD     = dispatch_request
    do_POST     = dispatch_request
    do_PUT      = dispatch_request
    do_DELETE   = dispatch_request
    do_GET      = dispatch_request
    

class JJProxyHTTPServer (SocketServer.ThreadingMixIn,
                           BaseHTTPServer.HTTPServer): pass

if __name__ == '__main__':
    from sys import argv
    if argv[1:] and argv[1] in ('-h', '--help'):
        print argv[0], "[port]"
    else:
        argv[1] = "8001"
        BaseHTTPServer.test(ProxyRequestDispatchor, JJProxyHTTPServer)
